/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file edge32_hor_c3_sum.cpp
 */
#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t LENGTH = 14586;
class KernelEdge32HorC3Sum {
public:
    __aicore__ inline KernelEdge32HorC3Sum() {}
    __aicore__ inline void Init(GM_ADDR img, GM_ADDR out,
                                uint32_t singleHeigh1, uint32_t width, uint32_t heigh,
                                TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        if (GetBlockIdx() < 8) 
        {
            this->singleHeigh = 262;

            imgGm.SetGlobalBuffer((__gm__ uint8_t*)img + this->singleHeigh * GetBlockIdx() * width * 3, this->singleHeigh * width * 3);
            outGm.SetGlobalBuffer((__gm__ uint8_t*)out + this->singleHeigh * GetBlockIdx() * width * 3, this->singleHeigh * width * 3);
        }
        else
        {
            this->singleHeigh = 192;

            imgGm.SetGlobalBuffer((__gm__ uint8_t*)img + 262*8*width * 3 + this->singleHeigh * (GetBlockIdx()-8) * width * 3, this->singleHeigh * width * 3);
            outGm.SetGlobalBuffer((__gm__ uint8_t*)out + 262*8*width * 3 + this->singleHeigh * (GetBlockIdx()-8) * width * 3, this->singleHeigh * width * 3);
        }

        this->width = width;
        this->heigh = heigh;
        this->tileDataNum = 4896*3;

        pipeIn->InitBuffer(inQueueImg, BUFFER_NUM, this->tileDataNum * 4 * sizeof(uint8_t));
        pipeIn->InitBuffer(outQueueOut, BUFFER_NUM, this->tileDataNum * sizeof(uint8_t));

        pipeIn->InitBuffer(QueueTmp1, this->tileDataNum * sizeof(float));
        pipeIn->InitBuffer(QueueTmp2, this->tileDataNum * sizeof(half));
    }
    __aicore__ inline void Process() {
        auto tmp1 = QueueTmp1.Get<float>();
        int32_t loopCount = this->singleHeigh;
        SetDeqScale((half)1);
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i, 0);
            Compute1(i, tmp1);
            CopyIn(i, 4);
            Compute2(i, tmp1);
            CopyOut(i);
        }
    }
private:
    __aicore__ inline void CopyIn(int32_t progress, int32_t offset) {
        LocalTensor<uint8_t> imgLocal = inQueueImg.AllocTensor<uint8_t>();
        DataCopy(imgLocal, imgGm[(progress * this->width+9 + offset) * 3], this->tileDataNum);
        DataCopy(imgLocal[this->tileDataNum], imgGm[(progress * this->width+9 + offset+1) * 3], this->tileDataNum);
        DataCopy(imgLocal[this->tileDataNum*2], imgGm[(progress * this->width+9 + offset+2) * 3], this->tileDataNum);
        DataCopy(imgLocal[this->tileDataNum*3], imgGm[(progress * this->width+9 + offset+3) * 3], this->tileDataNum);
        inQueueImg.EnQue(imgLocal);
    }
    __aicore__ inline void Compute1(int32_t progress, LocalTensor<float> tmp1) {
        LocalTensor<uint8_t> imgLocal = inQueueImg.DeQue<uint8_t>();
        auto tmp2 = QueueTmp2.Get<half>();
        auto tmp3 = imgLocal.ReinterpretCast<half>();
        Cast(tmp2, imgLocal, RoundMode::CAST_NONE, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum], RoundMode::CAST_NONE, this->tileDataNum);
        Add(tmp2, tmp2, tmp3, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum*2], RoundMode::CAST_NONE, this->tileDataNum);
        Add(tmp2, tmp2, tmp3, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum*3], RoundMode::CAST_NONE, this->tileDataNum);

        Add(tmp1.ReinterpretCast<half>(), tmp2, tmp3, this->tileDataNum);
        inQueueImg.FreeTensor(imgLocal);
    }
    __aicore__ inline void Compute2(int32_t progress, LocalTensor<float> tmp1) {
        LocalTensor<uint8_t> imgLocal = inQueueImg.DeQue<uint8_t>();
        LocalTensor<uint8_t> outLocal = outQueueOut.AllocTensor<uint8_t>();
        auto tmp2 = QueueTmp2.Get<half>();
        auto tmp3 = imgLocal.ReinterpretCast<half>();
        auto tmp3_f = imgLocal.ReinterpretCast<float>();

        Duplicate(outLocal.ReinterpretCast<half>()[7280], (half)0, 48);

        Cast(tmp2, imgLocal, RoundMode::CAST_NONE, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum], RoundMode::CAST_NONE, this->tileDataNum);
        Add(tmp2, tmp2, tmp3, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum*2], RoundMode::CAST_NONE, this->tileDataNum);
        Add(tmp2, tmp2, tmp3, this->tileDataNum);
        Cast(tmp3, imgLocal[this->tileDataNum*3], RoundMode::CAST_NONE, this->tileDataNum);
        Add(tmp2, tmp2, tmp3, this->tileDataNum);

        Add(tmp2, tmp1.ReinterpretCast<half>(), tmp2, this->tileDataNum);
        Cast(tmp1, tmp2, RoundMode::CAST_NONE, this->tileDataNum);

        Add(tmp1, tmp1, tmp1[24], LENGTH);
        Muls(tmp1, tmp1, (float)0.0625, LENGTH);

        Cast(tmp3_f.ReinterpretCast<int32_t>(), tmp1, RoundMode::CAST_TRUNC, LENGTH);
        Cast(tmp1.ReinterpretCast<half>(), tmp3_f.ReinterpretCast<int32_t>(), RoundMode::CAST_NONE, LENGTH);
        Cast(outLocal, tmp1.ReinterpretCast<half>(), RoundMode::CAST_NONE, LENGTH);

        inQueueImg.FreeTensor(imgLocal);
        outQueueOut.EnQue<uint8_t>(outLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<uint8_t> outLocal = outQueueOut.DeQue<uint8_t>();
        DataCopy(outGm[(progress * this->width + 16) * 3], outLocal, 14656);
        outQueueOut.FreeTensor(outLocal);
    }

private:
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueImg;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;

    GlobalTensor<uint8_t> imgGm;
    GlobalTensor<uint8_t> outGm;

    uint32_t singleHeigh;
    uint32_t tileDataNum;

    uint32_t width;
    uint32_t heigh;
};
extern "C" __global__ __aicore__ void edge32_hor_c3_sum(GM_ADDR img, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    KernelEdge32HorC3Sum op;
    op.Init(img, out,
            tiling_data.singleHeigh, tiling_data.width, tiling_data.heigh, 
            &pipe);  
    KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_VECTOR_CORE);
    op.Process();
}